#include "ts_stream_demuxer.h"
#include <stdio.h>

using namespace stream_type;
using namespace tsde_status;


tsDemuxer::tsDemuxer()
{
	m_channel = 0;
}

tsDemuxer::~tsDemuxer()
{
}

status tsDemuxer::demuxTsPackets(const char* ts_data,int ts_len,uint8_t** frame_out,int *frame_len,uint32_t *timestamp,stype *type)
{
	int pkg_num = ts_len / 188;
	status res;
	const char * ptr = ts_data;

	while(pkg_num--)
	{
		res = demuxATsPacket(ptr,frame_out,frame_len,timestamp,type);
		ptr += 188;
	}

	return res;
}

status tsDemuxer::demuxATsPacket(const char* ts_data,uint8_t** frame_out,int *frame_len,uint32_t *timestamp,stype *type)
{
	status res = status_demux_more_data;
	*frame_out = 0;
	*frame_len = 0;
	*timestamp = 0;
	*type = data;
	const char* ptr = ts_data;
	const char* end_ptr = ptr+188;

	if(ptr[0] != 0x47)            // ts sync byte
		return status_demux_err;

	uint16_t pid = toInt(ptr+1);
	uint8_t flags = toByte(ptr+3);

	bool transport_error = pid & 0x8000;
	bool payload_unit_start_indicator = pid & 0x4000;
	bool adaptation_field_exist = flags & 0x20;
	bool payload_data_exist = flags & 0x10;
	//uint8_t continuity_counter = flags & 0x0f;
	pid&=0x1fff;

	if(transport_error)
		return status_demux_err;

	if(pid == 0x1fff || !payload_data_exist)
		return status_demux_unkonw;  //0

	ptr += 4;

	// skip adaptation field
	if(adaptation_field_exist)
	{
		ptr += toByte(ptr) + 1;
		if(ptr >= end_ptr)
			return status_demux_err;
	}

	stream& s = m_streams[pid];

	if(!pid || (s.channel != 0xffff && s.type == 0xff))
	{
		// PSI
	    if(payload_unit_start_indicator)
	    {
	    	// begin of PSI table
	    	ptr++;

	    	if(ptr >= end_ptr)
	    		return status_demux_err;

	    	if(*ptr != 0x00 && *ptr != 0x02)
	    		return status_demux_unkonw;  //0

	    	if(end_ptr - ptr < 3)
	    		return status_demux_err;

	    	uint16_t l = toInt(ptr + 1);

	    	if(l&0x3000 != 0x3000)
	    		return status_demux_err;

	    	l &= 0x0fff;
	    	ptr += 3;

	    	int len = end_ptr-ptr;

	    	if(l > len)
	    	{
	    		if(l > table::max_buf_len)
	    			return status_demux_err;

	    		s.psi.reset();

	    		memcpy(s.psi.buf,ptr,len);
	    		s.psi.offset += len;
	    		s.psi.len = l;

	    		return status_demux_unkonw; //0
	    	}
	    	else
	    		end_ptr=ptr + l;
	    }
	    else
	    {
	    	// next part of PSI
	    	if(!s.psi.offset)
	    		return status_demux_err;

	    	int len = end_ptr - ptr;

	    	if(len > table::max_buf_len - s.psi.offset)
	    		return status_demux_err;

	    	memcpy(s.psi.buf + s.psi.offset,ptr,len);
	    	s.psi.offset += len;

	    	if(s.psi.offset < s.psi.len)
	    		return status_demux_unkonw;  //0
	    	else
	    	{
	    		ptr = s.psi.buf;
	    		end_ptr = ptr+s.psi.len;
	    	}
	    }

        if(!pid)
        {
        	// PAT
            ptr += 5;
            if(ptr >= end_ptr)
                return status_demux_err;

            int len = end_ptr - ptr - 4;
            if(len < 0 || len%4)
                return status_demux_err;

            int n = len / 4;

            for(int i = 0;i < n;i++,ptr += 4)
            {
                uint16_t channel = toInt(ptr);
                uint16_t pid = toInt(ptr+2);

                if(pid & 0xe000 != 0xe000)
                    return status_demux_err;

                pid &= 0x1fff;

                if(!m_channel || m_channel == channel)
                {
                    stream& ss = m_streams[pid];
                    ss.channel = channel;
                    ss.type = 0xff;
                }
            }
        }
        else
        {
            // PMT

            ptr += 7;

            if(ptr >= end_ptr)
                return status_demux_err;

            uint16_t info_len = toInt(ptr)&0x0fff;

            ptr += info_len+2;
            end_ptr -= 4;

            if(ptr >= end_ptr)
                return status_demux_err;

            while(ptr < end_ptr)
            {
                if(end_ptr - ptr < 5)
                    return status_demux_err;

                uint8_t type = toByte(ptr);
                uint16_t pid = toInt(ptr+1);

                if(pid & 0xe000 != 0xe000)
                    return status_demux_err;

                pid &= 0x1fff;

                info_len = toInt(ptr+3)&0x0fff;

                ptr += 5 + info_len;

                // ignore unknown streams
                if(validateType(type))
                {
                    stream& ss = m_streams[pid];

                    if(ss.channel != s.channel || ss.type != type)
                    {
                        ss.channel = s.channel;
                        ss.type = type;
                        ss.id = ++s.id;

                        if(getStreamType(type) <= vc1_video)  //video
                        {
                        	ss.buffer[0] = new uint8_t[VIODE_BUFFER_SIZE];
                        	ss.buffer[1] = new uint8_t[VIODE_BUFFER_SIZE];
                        }
                        else                                               //audio
                        {
                        	ss.buffer[0] = new uint8_t[AUDIO_BUFFER_SIZE];
                        	ss.buffer[1] = new uint8_t[AUDIO_BUFFER_SIZE];
                        }
                    }
                }
            }

            if(ptr != end_ptr)
                return status_demux_err;
        }
	}
	else
	{
		if(s.type != 0xff)
		{
			// PES

			if(payload_unit_start_indicator)
			{
				s.psi.reset();
				s.psi.len = 9;
			}

			while(s.psi.offset < s.psi.len)
			{
				int len = end_ptr-ptr;

				if(len <= 0)
					return status_demux_unkonw; //0

				int n = s.psi.len - s.psi.offset;

				if(len > n)
					len = n;

				memcpy(s.psi.buf + s.psi.offset,ptr,len);
				s.psi.offset += len;

				ptr += len;

				if(s.psi.len == 9)
					s.psi.len += toByte(s.psi.buf+8);
			}

			if(s.psi.len)
			{
				if(memcmp(s.psi.buf,"\x00\x00\x01",3))
					return status_demux_err;

				s.stream_id = toByte(s.psi.buf+3);
				uint8_t flags = toByte(s.psi.buf+7);

				if(s.frame_num > 0)
				{
					*frame_out = s.buffer[s.frame_num%2];
					*frame_len = s.frame_len;
					*timestamp = s.timestamp;
					*type = (stype)getStreamType(s.type);
					res = status_demux_ok;

				}

				s.frame_len = 0;
				s.frame_num++;

				switch(flags&0xc0)
				{
				case 0x80:          // PTS only
				{
					uint64_t pts = decodePts(s.psi.buf+9);

					s.timestamp = pts/90;

					if(s.dts > 0 && pts > s.dts)
						s.frame_length = pts-s.dts;
					s.dts = pts;

					if(pts > s.last_pts)
						s.last_pts = pts;

					if(!s.first_pts)
						s.first_pts = pts;
				}
				break;
				case 0xc0:          // PTS,DTS
				{
					uint64_t pts = decodePts(s.psi.buf+9);
					uint64_t dts = decodePts(s.psi.buf+14);

					s.timestamp = pts/90;

					if(s.dts > 0 && dts > s.dts)
						s.frame_length = dts - s.dts;
					s.dts = dts;

					if(pts > s.last_pts)
						s.last_pts = pts;

					if(!s.first_dts)
						s.first_dts = dts;
				}
				break;
				}
				s.psi.reset();
			}

			if(s.frame_num)
			{
				int len = end_ptr-ptr;
				memcpy(s.buffer[s.frame_num%2] + s.frame_len,ptr,len);
				s.frame_len += len;
			}
		}
	}
	return res;
}

inline bool tsDemuxer::validateType(uint8_t type)
{
	return strchr("\x01\x02\x80\x1b\xea\x81\x06\x83\x03\x04\x82\x86\x8a\x11\x0f",type)?true:false;
}

uint64_t tsDemuxer::decodePts(const char* ptr)
{
    const unsigned char* p=(const unsigned char*)ptr;

    uint64_t pts=((p[0]&0xe)<<29);
    pts|=((p[1]&0xff)<<22);
    pts|=((p[2]&0xfe)<<14);
    pts|=((p[3]&0xff)<<7);
    pts|=((p[4]&0xfe)>>1);

    return pts;
}

int tsDemuxer::getStreamType(uint8_t type)
{
    switch(type)
    {
    case 0x01:
    case 0x02:
        return mpeg2_video;
    case 0x80:
        return mpeg2_video;
    case 0x1b:
        return h264_video;
    case 0xea:
        return vc1_video;
    case 0x81:
    case 0x06:
    case 0x83:
        return ac3_audio;
    case 0x03:
    case 0x04:
        return mpeg2_audio;
    case 0x82:
    case 0x86:
    case 0x8a:
        return dts_audio;
	case 0x11:
	case 0x0f:
		return aac_audio;

    }
    return data;
}

const char* tsDemuxer::getStreamExt(uint8_t type_id)
{
    static const char* list[9]= { "sup", "m2v", "264", "vc1", "ac3", "m2a", "pcm", "dts","aac" };

    if(type_id<0 || type_id>=9)
        type_id=0;

    return list[type_id];
}

inline uint8_t tsDemuxer::toByte(const char* p)
{
	return *((unsigned char*)p);
}

inline uint16_t tsDemuxer::toInt(const char* p)
{
	uint16_t n=((unsigned char*)p)[0]; n<<=8; n+=((unsigned char*)p)[1];
	return n;
}

inline uint32_t tsDemuxer::toInt32(const char* p)
{
	uint32_t n=((unsigned char*)p)[0];
	n<<=8; n+=((unsigned char*)p)[1];
	n<<=8; n+=((unsigned char*)p)[2];
	n<<=8; n+=((unsigned char*)p)[3];
	return n;
}
